"""
修复版本的DeepPA模型，解决GPU FFT问题
"""
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))

# 导入原始模型
from src.models.DeepPA import DeepPA as OriginalDeepPA

class DeepPA(OriginalDeepPA):
    """
    修复版本的DeepPA，解决GPU FFT问题
    """
    
    def _gco_fft(self, x, bias):
        """
        修复的GCO FFT实现，强制使用CPU FFT避免GPU FFT错误
        """
        dtype = x.dtype
        x = x.float()
        B, N, C = x.shape

        # 强制使用CPU FFT避免GPU FFT错误
        _dev = bias.device
        x_cpu = x.to("cpu").contiguous()
        x = torch.fft.rfft(x_cpu, dim=1, norm="ortho").to(_dev)
        
        # 其余逻辑保持不变
        x = x * bias
        x = torch.fft.irfft(x, n=N, dim=1, norm="ortho")
        return x.to(dtype)

# 导出修复后的模型
__all__ = ['DeepPA']
